Mooncake: Kimi’s KVCache-centric Architecture for LLM Serving](https://arxiv.org/abs/2407.00079)
月之暗面kimi团队在2024的一篇工作
FAST 2025 Best Paper
Abstract
- 以KV cache为核心
- 采用了prefill和decode的分解架构
- 利用了集群中未被使用的CPU、DRAM和SSD等资源
- 在最大化吞吐量的同时,还平衡了服务级别目标
- 面临的是高度超载的环境,提出了基于预测的早期拒绝政策,
- Model as a Service (MaaS)的最大目标是 maximize overall effective throughput和尽量约束varying levels of SLOs,比如 the time to first token (TTFT)和 the time between tokens (TBT)。
- prefill和decoding计算特性不一样,前者是compute bound,后者是memory bound,曾有相关工作提出了分离架构。
以KV Cache为中心的架构
分为三部分:
- Prefill Pool
- max Cache Reuse
- Mooncake Store
- Decoding Pool
- max Throughput
Workfolw
KVCache Reuse
- 尽可能多地reuse KV Cache:一边计算A一边传输B,可以理解为减轻计算负载;同时还能减低TTFT时间
- 平衡各节点负载
- 满足TTFT
Incremental Prefill
- 以块的形式计算和传输KV Cache,减少KV Cache传输造成的阻塞时间
KVCache Transfer
- 和Prefill异步,从而实现overlap
Decoding
根据instance负载选择合适的decoding instance
当KV Cache传输到Decoding Instance完成时,加入下一iteration的计算
Mooncake Store
【3.2.1】KV Cache管理:存储的元数据,比如说前缀信息HashMap
【3.2.2】API接口介绍
【3.2.3】设计传输引擎
目标:
- 有效使用多个RDMA NIC
- 抽象出来的API
- 支持错误处理
NCCL不支持动态调整和DRAM to DRAM的路径
1. Topology-aware path selection
基于拓扑的链接
- 广播拓扑
- 建立链接
分块传输(16KB)
2. Endpoint Managerment
维护了一个连接的Endpoint Pool,按需生成Endpoint,Pool限制最大数量,采用了SIEVE算法。
3. Failure Handing
自动识别损坏的路径并更新拓扑
Mooncake‘s Prefill Pool
Chunked pipeline parallelism
Cache-aware的全局调度策略
- 基于请求长度、prefix cache命中率和local queuing time,计算各个新请求在不同Instance上的TTFT,分配给TTFT最小的
- TTFT最小的节点不一定是持有最大前缀匹配的instance,该instance可以主动地向持有其最大前缀匹配的节点获取KV Cache数据。
- 如果本地的前缀匹配token数乘以一个阈值大于所有节点的最大前缀匹配token数,直接本地计算。
跑了16个8*A100 nodes的实验证明,这种方式TTFT更优。
实验结果
在不同输入和更长的输出情况下(Chatbot),vLLM TBT违反率很高。
在较短的输入和输出时,vLLM多节点前缀利用率不如Mooncake。
在混合各种dataset的实验,Mooncake表现依旧更好。
Prefill时间实验。
利用global cache后的前缀缓存命中率实验:
cache-aware机制下缓存块副本的数量:比较稳定
通信实验:
100Gbps是Mooncake这种Cache机制下的底线
前缀缓存降低了很多时间,虽然其他部分时间增长,但前缀缓存还是有收益
PD实验:1比1 其实比较稳定,可以在不同时间段动态切换一下